import getpass

from langchain.agents import AgentExecutor
from langchain.agents.openai_assistant import OpenAIAssistantRunnable
from langchain.schema import AgentFinish
from langchain.tools.ddg_search import DuckDuckGoSearchRun
from langchain.tools.e2b_data_analysis.tool import E2BDataAnalysisTool
tools = [E2BDataAnalysisTool(api_key=getpass.getpass()), DuckDuckGoSearchRun()]
agent = OpenAIAssistantRunnable.create_assistant(
    name="langchain assistant e2b tool",
    instructions="You are a personal math tutor. Write and run code to answer math questions. You can also search the internet.",
    tools=tools,
    model="gpt-3.5-turbo",
    as_agent=True,
)


def init_agent_with_assistant():
    interpreter_assistant = OpenAIAssistantRunnable.create_assistant(
        name="langchain assistant",
        instructions="You are a personal math tutor. Write and run code to answer math questions.",
        tools=[{"type": "code_interpreter"}],
        model="gpt-3.5-turbo",
    )
    output = interpreter_assistant.invoke({"content": "What's 10 - 4 raised to the 2.7"})
    print(output)
    agent_executor = AgentExecutor(agent=agent, tools=tools)
    agent_executor.invoke({"content": "What's the weather in SF today divided by 2.7"})

    agent = OpenAIAssistantRunnable.create_assistant(
        name="langchain assistant e2b tool",
        instructions="You are a personal math tutor. Write and run code to answer math questions.",
        tools=tools,
        model="gpt-4-1106-preview",
        as_agent=True,
    )


def execute_agent(agent, tools, input):
    tool_map = {tool.name: tool for tool in tools}
    response = agent.invoke(input)
    while not isinstance(response, AgentFinish):
        tool_outputs = []
        for action in response:
            tool_output = tool_map[action.tool].invoke(action.tool_input)
            print(action.tool, action.tool_input, tool_output, end="\n\n")
            tool_outputs.append(
                {"output": tool_output, "tool_call_id": action.tool_call_id}
            )
        response = agent.invoke(
            {
                "tool_outputs": tool_outputs,
                "run_id": action.run_id,
                "thread_id": action.thread_id,
            }
        )

    return response


if __name__ == '__main__':
    response = execute_agent(agent, tools, {"content": "What's 10 - 4 raised to the 2.7"})
    print(response.return_values["output"])

    # 使用现有的助手
    agent = OpenAIAssistantRunnable(assistant_id="<ASSISTANT_ID>", as_agent=True)

